from openpyxl import load_workbook
from collections import defaultdict

# 0-based index
CFO_ID_COL_IDX = 2
COMP_ID_COL_IDX = 4
YEAR_COL_IDX = 5
GENDER_COL_IDX = 6
# 1-based index
MM_COL_IDX = 9
MF_COL_IDX = 10
FM_COL_IDX = 11
FF_COL_IDX = 12
POST_COL_IDX = 13


def find_split(l):
    for i in range(len(l)):
        before_cfo = l[i - 1]['cfo'] if i > 0 else None
        after_cfo = l[i]['cfo'] if i < len(l) else None
        if before_cfo == after_cfo:
            continue
        before_counting = 0
        after_counting = 0
        before_bounder = i - 1
        after_bounder = i + 1
        for j in range(i - 1, -1, -1):  # [i-1,0]
            if before_cfo == l[j]['cfo']:
                before_counting += 1
                before_cfo = l[j]['cfo']
                before_bounder = j
            else:
                break
        for k in range(i + 1, len(l), 1):  # [i+1, len(l)-1]
            if after_cfo == l[k]['cfo']:
                after_counting += 1
                after_cfo = l[k]['cfo']
                after_bounder = k
            else:
                break
        if before_counting >= 3 and after_counting >= 3:
            changing_type = l[i - 1]['gender'].capitalize() + '_to_' + l[i]['gender'].capitalize()
            for a in range(before_bounder, after_bounder + 1):
                l[a][changing_type] = 1
            for a in range(before_bounder, i):
                l[a]['Post'] = 0
            for a in range(i+1, after_bounder+1):
                l[a]['Post'] = 1
            l[i]['Post'] = 'None'
            break
    return l



def main():
    # filename = 'dataset/V4.xlsx'
    filename = 'dataset/CFO Change 1992-2019-V4.xlsx'
    out_file = 'result/CFO Change 1992-2019-V4_with_result.xlsx'
    # out_file = 'result/V4_with_result.xlsx'

    wb = load_workbook(filename=filename)
    ws = wb['Final']
    cnt = 0
    group_records = defaultdict(list)
    is_header = True
    for i, r in enumerate(ws.rows):
        if is_header:
            is_header = False
            continue
        cfo = r[CFO_ID_COL_IDX].value  # starting from 0
        comp = r[COMP_ID_COL_IDX].value
        year = r[YEAR_COL_IDX].value
        gender = r[GENDER_COL_IDX].value
        group_records[comp].append({
            'rn': i+1,
            'cfo': cfo,
            'comp': comp,
            'year': year,
            'gender': gender
        })
        cnt += 1
        # if cnt > 1000:
        #     break

    print("load done")
    for k, v in group_records.items():
        find_split(v)
        # group_records[k] =find_split(v)

    # valid_list = [x for x in group_records.values()]
    for cfo_list in group_records.values():
        for cfo in cfo_list:
            row_no = cfo['rn']
            if 'Male_to_Male' in cfo.keys():
                ws.cell(row=row_no, column=MM_COL_IDX).value = 1
            if 'Male_to_Female' in cfo.keys():
                ws.cell(row=row_no, column=MF_COL_IDX).value = 1
            if 'Female_to_Male' in cfo.keys():
                ws.cell(row=row_no, column=FM_COL_IDX).value = 1
            if 'Female_to_Female' in cfo.keys():
                ws.cell(row=row_no, column=FF_COL_IDX).value = 1
            if 'Post' in cfo.keys():
                ws.cell(row=row_no, column=POST_COL_IDX).value = cfo['Post']


    # Save the file
    wb.save(out_file)


if __name__ == '__main__':
    main()
